import time, traceback, os
import numpy as np
import random

# 导入工具模块
from Utils.tensor_ops import my_view, __hash__, repeat_at, gather_righthand
from Common.action_set_encoder import encode_action_as_digits
from config import GlobalConfig
from collections import Counter


def str_array_to_num(str_arr):
    """
    将字符串数组转换为数字数组

    这个函数用于将重复的字符串映射为数字索引，便于神经网络处理。
    常用于将智能体类型、动作名称等字符串信息转换为数字表示。

    Args:
        str_arr (list): 字符串列表

    Returns:
        list: 对应的数字索引列表
    """
    out_arr = []
    buffer = {}  # 缓存字典，存储已见过的字符串到数字的映射
    for str in str_arr:
        if str not in buffer:
            buffer[str] = len(buffer)  # 为新字符串分配唯一的数字索引
        out_arr.append(buffer[str])
    return out_arr


def itemgetter(*items):
    """
    自定义的itemgetter函数，等同于operator.itemgetter

    这个函数用于从字典中提取指定的键值，返回一个元组。
    如果某个键不存在，则返回None而不是抛出异常。

    Args:
        *items: 要获取的键名列表

    Returns:
        function: 一个函数，接受字典并返回指定键的值元组
    """
    def g(obj):
        return tuple(obj[item] if item in obj else None for item in items)
    return g


class ActionConvertLegacy:
    """
    动作编码转换类：将RL模型生成的动作序列转化为约定的动作编号

    这个类负责将强化学习模型输出的动作索引转换为游戏引擎可以理解的动作格式。
    它还管理不同类型智能体的动作可用性，确保每个智能体只执行其类型的合法动作。
    """

    def __init__(self, SELF_TEAM_ASSUME, OPP_TEAM_ASSUME, OPP_NUM_ASSUME) -> None:
        """
        初始化智能体的动作空间

        Args:
            SELF_TEAM_ASSUME (int): 友方队伍编号（默认为0）
            OPP_TEAM_ASSUME (int): 敌方队伍编号（默认为1）
            OPP_NUM_ASSUME (int): 敌方单位数量
        """
        self.SELF_TEAM_ASSUME = SELF_TEAM_ASSUME
        self.OPP_TEAM_ASSUME = OPP_TEAM_ASSUME
        self.OPP_NUM_ASSUME = OPP_NUM_ASSUME

        # 生成随机整数，用于某些随机化操作（具体用途需要查看更多上下文）
        random_ints = [random.randint(1, 100) for _ in range(36)]

        # 定义动作字典：包含所有可能的动作组合
        # 每个条目格式：(main_cmd, sub_cmd, points, x, y, z, UID, T, T_index)
        self.dictionary_args = [
            # 待命动作
            ('Idle', 'N/A', None, None, None, None, None, None),   # 0 - 无动作

            # 警戒动作：8个方向
            ('Guard', '+X', None, None, None, None, None, None),  # 1 - 警戒+X方向
            ('Guard', '+Y', None, None, None, None, None, None),  # 2 - 警戒+Y方向
            ('Guard', '-X', None, None, None, None, None, None),  # 3 - 警戒-X方向
            ('Guard', '-Y', None, None, None, None, None, None),  # 4 - 警戒-Y方向
            ('Guard', '+X+Y', None, None, None, None, None, None),  # 5 - 警戒+X+Y方向
            ('Guard', '-X+Y', None, None, None, None, None, None),  # 6 - 警戒-X+Y方向
            ('Guard', '-X-Y', None, None, None, None, None, None),  # 7 - 警戒-X-Y方向
            ('Guard', '+X-Y', None, None, None, None, None, None),  # 8 - 警戒+X-Y方向

            # 移动动作：基础4个方向
            ('Moving', '+X', None, None, None, None, None, None),  # 9 - 移动+X方向
            ('Moving', '+Y', None, None, None, None, None, None),  # 10 - 移动+Y方向
            ('Moving', '-X', None, None, None, None, None, None),  # 11 - 移动-X方向
            ('Moving', '-Y', None, None, None, None, None, None),  # 12 - 移动-Y方向

            # 移动动作：对角线4个方向
            ('Moving', '+X+Y', None, None, None, None, None, None),  # 13 - 移动+X+Y方向
            ('Moving', '-X+Y', None, None, None, None, None, None),  # 14 - 移动-X+Y方向
            ('Moving', '-X-Y', None, None, None, None, None, None),  # 15 - 移动-X-Y方向
            ('Moving', '+X-Y', None, None, None, None, None, None),  # 16 - 移动+X-Y方向

            # 移动动作：包含Z轴的8个3D方向
            ('Moving', '+X+Z', None, None, None, None, None, None),  # 17 - 移动+X+Z方向
            ('Moving', '+Y+Z', None, None, None, None, None, None),  # 18 - 移动+Y+Z方向
            ('Moving', '-X+Z', None, None, None, None, None, None),  # 19 - 移动-X+Z方向
            ('Moving', '-Y+Z', None, None, None, None, None, None),  # 20 - 移动-Y+Z方向
            ('Moving', '+X+Y+Z', None, None, None, None, None, None),  # 21 - 移动+X+Y+Z方向
            ('Moving', '-X+Y+Z', None, None, None, None, None, None),  # 22 - 移动-X+Y+Z方向
            ('Moving', '-X-Y+Z', None, None, None, None, None, None),  # 23 - 移动-X-Y+Z方向
            ('Moving', '+X-Y+Z', None, None, None, None, None, None),  # 24 - 移动+X-Y+Z方向

            # 移动动作：负Z轴的8个3D方向
            ('Moving', '+X-Z', None, None, None, None, None, None),  # 25 - 移动+X-Z方向
            ('Moving', '+Y-Z', None, None, None, None, None, None),  # 26 - 移动+Y-Z方向
            ('Moving', '-X-Z', None, None, None, None, None, None),  # 27 - 移动-X-Z方向
            ('Moving', '-Y-Z', None, None, None, None, None, None),  # 28 - 移动-Y-Z方向
            ('Moving', '+X+Y-Z', None, None, None, None, None, None),  # 29 - 移动+X+Y-Z方向
            ('Moving', '-X+Y-Z', None, None, None, None, None, None),  # 30 - 移动-X+Y-Z方向
            ('Moving', '-X-Y-Z', None, None, None, None, None, None),  # 31 - 移动-X-Y-Z方向
            ('Moving', '+X-Y-Z', None, None, None, None, None, None),  # 32 - 移动+X-Y-Z方向

            # 移动动作：纯Z轴方向
            ('Moving', '+Z', None, None, None, None, None, None),  # 33 - 移动+Z方向
            ('Moving', '-Z', None, None, None, None, None, None),  # 34 - 移动-Z方向

            # 手榴弹攻击动作：3个方向
            ('GrenadeAttacking', '+X+Y', None, None, None, None, None, None),  # 35 - 手榴弹攻击+X+Y方向
            ('GrenadeAttacking', '-X+Y', None, None, None, None, None, None),  # 36 - 手榴弹攻击-X+Y方向
            ('GrenadeAttacking', '+Y', None, None, None, None, None, None)    # 37 - 手榴弹攻击+Y方向
        ]

        # 如果存在敌方智能体，为每个敌方智能体添加攻击动作
        if self.OPP_NUM_ASSUME > 0:
            for i in range(self.OPP_NUM_ASSUME):
                # 为每个敌方智能体创建普通攻击动作
                self.dictionary_args.append(('NormalAttacking', 'N/A', None, None, None, None, OPP_TEAM_ASSUME, i))

    def convert_act_arr(self, type, a):
        """
        修改无人机的可执行动作，比如无人机不执行攻击，无人机只进行静态警戒等

        根据智能体类型过滤和转换动作，确保每个智能体只执行其类型允许的动作。
        例如，某些类型的智能体可能无法执行攻击动作。

        Args:
            type (str): 智能体类型（如 'Drone', 'Vehicle', 'BP_RoboDog_C' 等）
            a (int): 动作在动作列表中的索引

        Returns:
            numpy.ndarray: 编码后的动作数组，维度为[8]
        """
        # 获取动作参数
        args = self.dictionary_args[a]

        # 特定类型智能体的动作限制
        if type == 'BP_RoboDog_C' or type == 'BP_MNWS_Vehicle_6x6UGV_C' or type == 'BP_Base_UAV_C':
            # 机器狗、6x6无人车、基础无人机不能执行手榴弹攻击
            if args[0] == 'GrenadeAttacking':
                # 将手榴弹攻击替换为待命动作
                return encode_action_as_digits('Idle', 'N/A', None, None, None, None, None, None)
            # 其他动作正常执行
            return encode_action_as_digits(*args)
        else:
            # 其他类型智能体可以执行所有动作
            return encode_action_as_digits(*args)

    def get_tp_avail_act(self, type, uid):
        """
        生成智能体可执行的动作标志序列

        该序列的作用是根据用户的配置，确定是否对RL模型生成的动作进行校正，
        以屏蔽一些不合理的动作。这是一个动作掩码机制。

        Args:
            type (str): 智能体类型（如无人机、士兵等）
            uid (int): 智能体的全局唯一ID

        Returns:
            numpy.ndarray: 0和1组成的数组，shape(动作数量)
                          1表示动作可执行，0表示动作不可执行
        """
        DISABLE = 0  # 动作不可用
        ENABLE = 1   # 动作可用
        n_act = len(self.dictionary_args)  # 总动作数
        ret = np.zeros(n_act) + ENABLE     # 初始化所有动作为可用

        # 遍历所有动作，根据智能体类型和当前状态设置动作可用性
        for i in range(n_act):
            args = self.dictionary_args[i]

            # 全局限制：对所有类型智能体的限制
            # 例如，可以在这里禁止某些动作类型
            # if args[0] == 'PatrolMoving': ret[i] = DISABLE

            # 特定类型智能体的动作限制
            if type == 'BP_RoboDog_C' or type == 'BP_MNWS_Vehicle_6x6UGV_C' or type == 'BP_Base_UAV_C':
                # 机器狗、6x6无人车、基础无人机不能执行手榴弹攻击
                if args[0] == 'GrenadeAttacking':
                    ret[i] = DISABLE

            if type == 'BP_RoboDog_C' or type == 'BP_MNWS_Vehicle_6x6UGV_C':
                # 机器狗和6x6无人车不能进行垂直移动（±Z方向）
                if args[0] == 'Moving':
                    if args[1] == '+Z' or args[1] == '-Z':
                        ret[i] = DISABLE

            if type == 'BP_MNWS_Vehicle_6x6UGV_C':
                # 6x6无人车不能执行警戒动作
                if args[0] == 'Guard':
                    ret[i] = DISABLE

            # 自我攻击检查：智能体不能攻击自己
            if (args[0] == 'NormalAttacking' or args[0] == 'GrenadeAttacking') and args[7] in uid:
                ret[i] = DISABLE

        return ret


class RLAlgorithmBase:
    """
    强化学习算法基类

    这是一个抽象基类，定义了所有强化学习算法必须实现的基本接口。
    它处理与环境交互、决策制定、数据收集等通用功能。
    """

    def __init__(self, n_agent, n_thread, space, team=None):
        """
        初始化强化学习算法

        Args:
            n_agent (int): 智能体数量
            n_thread (int): 并行线程数量
            space (dict): 空间配置，包含观察空间和动作空间
            team (int, optional): 队伍编号
        """
        # 基本配置
        self.n_thread = n_thread           # 并行线程数
        self.n_agent = n_agent             # 智能体数量
        self.team = team                   # 队伍编号
        self.act_space = space['act_space']  # 动作空间
        self.obs_space = space['obs_space']  # 观察空间
        self.ScenarioConfig = GlobalConfig.ScenarioConfig  # 场景配置
        self.device = GlobalConfig.device   # 计算设备（CPU/GPU）

        # 内部状态变量
        self.avail_act = None              # 可用动作掩码
        self.opp_agent_not_alive = None    # 敌方智能体存活状态
        self._unfi_frag_ = None            # 未完成的轨迹片段
        self.trajectory_pool = None        # 轨迹池

        # 根据队伍设置日志目录
        if self.team == 0:
            self.logdir = GlobalConfig.logdir_red    # 红队日志目录
        elif self.team == 1:
            self.logdir = GlobalConfig.logdir_blue   # 蓝队日志目录
        else:
            print('[Save Model]: Error model path!')

        # 根据队伍设置智能体数量
        if self.team == 0:
            GlobalConfig.AlgorithmConfig_Red.num_agent = n_agent
        else:
            GlobalConfig.AlgorithmConfig_Blue.num_agent = n_agent

        # 奖励共享机制配置
        self.Enable_Reward_Sharing = False
        if hasattr(GlobalConfig.ScenarioConfig, 'Enable_Reward_Sharing'):
            self.Enable_Reward_Sharing = GlobalConfig.ScenarioConfig.Enable_Reward_Sharing

        # 观察空间配置
        if self.ScenarioConfig.EntityOriented:
            self.rawob_dim = self.ScenarioConfig.obs_vec_length  # 实体导向的观察维度
        else:
            self.rawob_dim = space['obs_space']['obs_shape']    # 传统观察空间维度

        # 初始化动作转换器
        if len(GlobalConfig.ScenarioConfig.N_AGENT_EACH_TEAM) >= 2:
            self.action_converter = ActionConvertLegacy(
                SELF_TEAM_ASSUME=team,
                OPP_TEAM_ASSUME=(1 - team),
                OPP_NUM_ASSUME=GlobalConfig.ScenarioConfig.N_AGENT_EACH_TEAM[1 - team]
            )
        else:
            self.action_converter = ActionConvertLegacy(
                SELF_TEAM_ASSUME=team,
                OPP_TEAM_ASSUME=1,
                OPP_NUM_ASSUME=0
            )

        # 动作数量
        self.n_actions = len(self.action_converter.dictionary_args)

        # 调试和检查参数
        self.patience = 2000  # 耐心值，用于调试检查

    def set_algorithm_trainer(self):
        """设置算法训练器 - 需要在子类中实现"""
        pass

    def interact_with_env(self, StateRecall):
        """
        与环境交互的主要方法

        这个方法是算法的核心，负责：
        1. 接收环境状态
        2. 进行决策
        3. 返回动作给环境
        4. 处理训练数据

        Args:
            StateRecall (dict): 环境状态回忆，包含观察、奖励、完成标志等

        Returns:
            tuple: (actions_list, StateRecall) - 动作列表和更新后的状态
        """
        # 获取智能体类型信息（如果尚未获取）
        if not hasattr(self, 'agent_type'):
            self.agent_uid = GlobalConfig.ScenarioConfig.AGENT_ID_EACH_TEAM[self.team]
            self.agent_type = [agent['type'] for agent in GlobalConfig.ScenarioConfig.SubTaskConfig.agent_list
                               if agent['team'] == self.team]

        # 获取敌方智能体存活状态
        if self.team == 0:
            self.opp_agent_not_alive = StateRecall['opp-agent-not-alive_team1']
        else:
            self.opp_agent_not_alive = StateRecall['opp-agent-not-alive_team0']

        # 处理动作过滤（根据队伍配置）
        if self.team == 0 and GlobalConfig.AlgorithmConfig_Red.action_filter_enabled:
            self.avail_act = np.array([np.stack(
                tuple(self.action_converter.get_tp_avail_act(tp, self.opp_agent_not_alive[index])
                      for tp in self.agent_type))
                for index in range(self.opp_agent_not_alive.shape[0])])

        if self.team == 1 and GlobalConfig.AlgorithmConfig_Blue.action_filter_enabled:
            self.avail_act = np.array([np.stack(
                tuple(self.action_converter.get_tp_avail_act(tp, self.opp_agent_not_alive[index])
                      for tp in self.agent_type))
                for index in range(self.opp_agent_not_alive.shape[0])])

        # 处理观察数据
        obs = StateRecall['Latest-Obs']
        obs = my_view(obs, [0, 0, -1, self.rawob_dim])  # 调整观察数据形状
        obs[(obs == 0).all(-1)] = np.nan  # 将全零观察替换为NaN

        # 获取环境状态标志
        P = StateRecall['ENV-PAUSE']  # 环境暂停标志
        R = ~P                        # 环境运行标志
        RST = StateRecall['Env-Suffered-Reset']  # 环境重置标志

        # 处理回合重置
        if RST.all():
            EpRsn = np.random.rand(self.n_thread) < 0
            StateRecall['_EpRsn_'] = EpRsn

        # 准备算法输入状态
        obs_feed = obs[R]  # 只获取运行中线程的观察
        I_StateRecall = {
            'obs': obs_feed,
            'Test-Flag': StateRecall['Test-Flag'],
            '_EpRsn_': StateRecall['_EpRsn_'][R],
            'threads_active_flag': R,
            'Latest-Team-Info': StateRecall['Latest-Team-Info'][R],
        }

        # 添加可用动作信息（如果启用了动作过滤）
        if self.team == 0 and GlobalConfig.AlgorithmConfig_Red.action_filter_enabled:
            I_StateRecall.update({'avail_act': self.avail_act[R]})

        if self.team == 1 and GlobalConfig.AlgorithmConfig_Blue.action_filter_enabled:
            I_StateRecall.update({'avail_act': self.avail_act[R]})

        # 初始化动作数组
        act = np.zeros(shape=(self.n_thread, self.n_agent), dtype=int) - 1

        # 进行决策
        act_active, internal_recall = self.making_decision(I_StateRecall, StateRecall['Test-Flag'])
        act[R] = act_active  # 只设置运行中线程的动作

        # 调试检查：确保选择的动作都是可用的
        if self.team == 0 and GlobalConfig.AlgorithmConfig_Red.action_filter_enabled and self.patience > 0:
            self.patience -= 1
            assert (gather_righthand(self.avail_act, repeat_at(act, -1, 1), check=False)[R] == 1).all()

        if self.team == 1 and GlobalConfig.AlgorithmConfig_Blue.action_filter_enabled and self.patience > 0:
            self.patience -= 1
            assert (gather_righthand(self.avail_act, repeat_at(act, -1, 1), check=False)[R] == 1).all()

        # 将动作索引转换为编码后的动作数组
        actions_list = [[self.action_converter.convert_act_arr(self.agent_type[agentid], act)
                        for agentid, act in enumerate(th)] for th in act]

        # 注册回调函数（用于接收奖励和下一个状态）
        if not StateRecall['Test-Flag']:
            StateRecall['_hook_'] = internal_recall['_hook_']
            assert StateRecall['_hook_'] is not None

        return actions_list, StateRecall

    def making_decision(self, StateRecall, test_mode):
        """
        决策方法 - 需要在子类中实现

        Args:
            StateRecall (dict): 算法输入状态
            test_mode (bool): 是否为测试模式

        Returns:
            tuple: (actions, internal_recall) - 动作和内部回忆
        """
        raise NotImplementedError

    def save_model(self, update_cnt, info=None):
        """
        保存模型 - 需要在子类中实现

        Args:
            update_cnt (int): 更新次数
            info (str, optional): 附加信息
        """
        raise NotImplementedError

    def process_framedata(self, traj_framedata):
        """
        处理轨迹帧数据

        这个方法处理从环境收集到的单帧数据，包括奖励、状态、动作等，
        并将其格式化为适合训练的格式。

        Args:
            traj_framedata (dict): 单帧轨迹数据

        Returns:
            dict: 处理后的轨迹数据
        """
        # 处理奖励共享机制
        if self.ScenarioConfig.Enable_Reward_Sharing:
            traj_framedata['reward'] = repeat_at(traj_framedata['reward'], insert_dim=-1, n_times=self.n_agent)

        # 更新数据字段名称
        traj_framedata['_DONE_'] = traj_framedata.pop('done')  # 完成标志
        traj_framedata['_TOBS_'] = traj_framedata.pop('terminal_state') if 'terminal_state' in traj_framedata else None

        # 屏蔽暂停环境的数据
        traj_framedata = self.mask_paused_env(traj_framedata)

        # 将数据喂入轨迹池
        self.trajectory_pool.feed_traj_framedata(traj_framedata)
        return traj_framedata

    def check_reward_type(self, AlgorithmConfig):
        """
        检查奖励类型配置的一致性

        Args:
            AlgorithmConfig: 算法配置
        """
        if self.ScenarioConfig.Enable_Reward_Sharing != AlgorithmConfig.TakeRewardAsUnity:
            assert self.ScenarioConfig.Enable_Reward_Sharing
            assert not AlgorithmConfig.TakeRewardAsUnity
            print('Warning, the scenario (Environment) provide `Enable_Reward_Sharing`, but AlgorithmConfig does not `Enable_Reward_Sharing` !')
            print('If you continue, team reward will be duplicated to serve as individual rewards, wait 3s to proceed...')
            time.sleep(3)

    def mask_paused_env(self, frag):
        """
        屏蔽暂停环境的数据

        Args:
            frag (dict): 轨迹片段

        Returns:
            dict: 屏蔽后的轨迹片段
        """
        running = ~frag['_SKIP_']  # 运行中的环境
        if running.all():
            return frag  # 如果所有环境都在运行，直接返回

        # 过滤掉暂停环境的数据
        for key in frag:
            if not key.startswith('_') and hasattr(frag[key], '__len__') and len(frag[key]) == self.n_thread:
                frag[key] = frag[key][running]
        return frag

    def on_notify(self, message, **kargs):
        """
        通知回调函数，在测试例程完成时调用

        Args:
            message: 通知消息
            **kargs: 其他参数
        """
        self.save_model(
            self.trajectory_pool.update_cnt,
            info=str(kargs)
        )

    def trace_update_callback(self, unfi_frag, req_hook=True):
        """
        轨迹更新回调函数，在接收到奖励时调用

        Args:
            unfi_frag (dict): 未完成的轨迹片段
            req_hook (bool): 是否需要钩子

        Returns:
            function or None: 钩子函数
        """
        assert self._unfi_frag_ is None
        self._unfi_frag_ = unfi_frag
        self._check_data_hash()  # 检查数据完整性
        if req_hook:
            return self.traj_waiting_hook
        else:
            return None

    def traj_waiting_hook(self, new_frag):
        """
        轨迹等待钩子函数

        这个函数将从 multi_team.py::deal_with_hook() 调用，
        当奖励和下一时刻观察准备好时被调用。

        Args:
            new_frag (dict): 新的轨迹片段
        """
        # 在开始时进行数据损坏检查
        self._check_data_curruption()

        # 用新数据完成帧数据
        fi_frag = self._unfi_frag_
        fi_frag.update(new_frag)

        # 调用上层函数处理帧数据
        self.process_framedata(traj_framedata=fi_frag)

        # 删除数据引用
        self._unfi_frag_ = None

    def _no_hook(self, new_frag):
        """无钩子函数"""
        return

    def _check_data_hash(self):
        """保护数据免被覆盖 - 检查数据哈希"""
        if self.patience > 0:
            self.patience -= 1
            self.hash_db = {}
            # 调试，检测写保护错误
            for key in self._unfi_frag_:
                item = self._unfi_frag_[key]
                if isinstance(item, dict):
                    self.hash_db[key] = {}
                    for subkey in item:
                        subitem = item[subkey]
                        self.hash_db[key][subkey] = __hash__(subitem)
                else:
                    self.hash_db[key] = __hash__(item)

    def _check_data_curruption(self):
        """保护数据免被覆盖 - 检查数据损坏"""
        if self.patience > 0:
            self.patience -= 1
            assert self._unfi_frag_ is not None
            assert self.hash_db is not None
            for key in self._unfi_frag_:
                item = self._unfi_frag_[key]
                if isinstance(item, dict):
                    for subkey in item:
                        subitem = item[subkey]
                        assert self.hash_db[key][subkey] == __hash__(subitem), ('Currupted data!')
                else:
                    assert self.hash_db[key] == __hash__(item), ('Currupted data!')

    def _create_config_fly(self):
        """创建配置文件 - 从子类移动到父类的方法"""
        logdir = GlobalConfig.logdir
        self.input_file_dir = '%s/cmd_io.txt' % logdir
        if not os.path.exists(self.input_file_dir):
            with open(self.input_file_dir, 'w+', encoding='utf8') as f:
                f.writelines(["# Write cmd at next line: ", ""])

    def _config_on_fly(self):
        """动态配置 - 从子类移动到父类的方法"""
        if not os.path.exists(self.input_file_dir):
            return

        with open(self.input_file_dir, 'r', encoding='utf8') as f:
            cmdlines = f.readlines()

        cmdlines_writeback = []
        any_change = False

        for cmdline in cmdlines:
            if cmdline.startswith('#') or cmdline == "\n" or cmdline == " \n":
                cmdlines_writeback.append(cmdline)
            else:
                any_change = True
                try:
                    print('[RLAgentAlgorithm.py] ------- executing: %s ------' % cmdline)
                    exec(cmdline)
                    cmdlines_writeback.append('# [execute successfully]\t' + cmdline)
                except:
                    print(traceback.format_exc())
                    cmdlines_writeback.append('# [execute failed]\t' + cmdline)

        if any_change:
            with open(self.input_file_dir, 'w+', encoding='utf8') as f:
                f.writelines(cmdlines_writeback)

    def generate_reward(self, state, env):
        """
        生成奖励函数

        Args:
            state: 当前状态
            env: 环境对象

        Returns:
            numpy.ndarray: 奖励数组
        """
        print("Call team-%d algorithm reward function" % self.team)
        reward = np.zeros(shape=self.n_agent)
        if self.Enable_Reward_Sharing:  # 如果启用奖励共享，每个队伍所有智能体共享奖励
            reward = np.zeros(shape=env.n_team)

        return reward